import sys, glob, os, re,subprocess
from collections import defaultdict
#from openpyxl import load_workbook, Workbook
from xlrd import open_workbook
from datetime import datetime
sys.path.append('/home/yong_sun/bin/workflow')
#from tools import ddRobot
#import requests

files = glob.glob(sys.argv[1] + '/*/*.kk.ra.xls') + glob.glob(f'/data/mNGS/runmngs/GZ/NCList/today/*.kk.ra.xls')

#date = sys.argv[1].split('/')[-1].split('_')[0]
date = datetime.strftime(datetime.today(),'%Y%m%d')

allsample = defaultdict(str)
sample = defaultdict(lambda: defaultdict(str))

samplenames = {}
nameid = []
species = defaultdict(str)
genus = defaultdict(str)

bodysitedict = defaultdict(str)
wbtmp = open_workbook(f'{sys.argv[4]}/body_site.xlsx')
sheet_names = wbtmp.sheet_names()
wstmp = wbtmp.sheet_by_name(sheet_names[0])
for i in range(1, wstmp.nrows):
    speciesname = wstmp.cell(i, 0).value
    speciesvalue = wstmp.cell(i, 1).value
    bodysitedict[speciesname] = speciesvalue

pathogendict = defaultdict(str)
wbtmp = open_workbook(f'{sys.argv[4]}/is_pathogen.xlsx')
sheet_names = wbtmp.sheet_names()
wstmp = wbtmp.sheet_by_name(sheet_names[0])
for i in range(1, wstmp.nrows):
    speciesname = wstmp.cell(i, 0).value
    speciesvalue = wstmp.cell(i, 1).value
    pathogendict[speciesname] = speciesvalue

genusdict = defaultdict(str)
wbtmp = open_workbook(f'{sys.argv[4]}/genus_zn_en.xlsx')
sheet_names = wbtmp.sheet_names()
wstmp = wbtmp.sheet_by_name(sheet_names[0])
for i in range(1, wstmp.nrows):
    speciesname = wstmp.cell(i, 0).value
    speciesvalue = wstmp.cell(i, 1).value
    genusdict[speciesname] = speciesvalue

ssize = {}
with open(sys.argv[3]) as ds:
    for d in ds:
        if d.strip():
            s = d.strip().split()
            ssize[s[0]] = f"{s[2]} "
        else:
            break

for i in files:
    # prefix = re.sub('.kk.ra.xls', '', i)
    name = i.split('/')[-1].split('.')
    if len(name) == 5:
        ssize[name[0]] = f"{int(name[1]):,}"
    if name[0] not in nameid:
        nameid.append(name[0])
    else:
        continue

    with open(i, encoding='gbk') as kk:
        kk.readline()
        for l in kk:
            e = l.split('\t')
            realspecies = ' '.join(e[3].split('_')[2:])
            realgenus = e[3].split('_')[2]
            for idx, cnt in enumerate(e[:5]):
                if not cnt:
                    e[idx] = '-'
            idstr = f'{e[0]}\t{e[2]}\t{e[3]}\t{e[4]}'
            allsample[idstr] = e[5]
            sample[idstr][name[0]] = e[5]
            species[idstr] = realspecies
            genus[idstr] = realgenus

stype = {}
with open(sys.argv[2]) as st:
    for i in st:
        e = i.strip().split()
        stype[e[0]] = e[2]

with open(sys.argv[1]+f'/mNGS.{date}.merge.KK.xls', 'w') as out:
    nameid.sort()
    namestr = []
    nametype = []
    namesize = []
    nc = []
    for i, e in enumerate(nameid[:]):
        if e.startswith('NC'):
            nc.append(e)
            nameid.remove(e)
    [nameid.insert(0, i) for i in reversed(nc)]
    for i, e in enumerate(nameid):
        namestr.append(e)
        if e in ssize:
            namesize.append(ssize[e])
        else:
            namesize.append('-')
        if e.find('-') > -1:
            n = e.split('-')[0]
        else:
            n = e
        if n in stype:
            nametype.append(stype[n])
        else:
            nametype.append('-')
    allsssdict = dict()
    out.write('-\t-\t-\t-\t' + '\t'.join(nametype)+'\t-\t-\t-\t-'+'\n')
    out.write('-\t-\t-\t-\t' + '\t'.join(namesize)+'\t-\t-\t-\t-'+'\n')
    out.write('D\tG\tS\tspecies_zn\t' + '\t'.join(namestr) +'\tsum\tgenus_zn\tbody_site\tis_pathogen' + '\n')

    for k, v in allsample.items():
        sss = '\t'.join([sample[k].get(i, '0') for i in nameid])  
        allsss = sum([int(i) for i in sss.split('\t')])
        allsssdict[f"{k}\t{sss}"] = str(allsss) + '\t' + genusdict[genus[k]] + '\t' + bodysitedict[species[k]] + '\t' +pathogendict[species[k]]
    allssslistsorted = sorted(allsssdict.items(), key = lambda x: x[1], reverse = True)
    for anytuple in allssslistsorted:
        out.write(f'{anytuple[0]}\t{str(anytuple[1])}\n')

from openpyxl import Workbook
from openpyxl.utils import get_column_letter
from openpyxl.styles import Font, PatternFill
from openpyxl.formatting.rule import CellIsRule
from openpyxl.styles import Border, borders, colors, Side
font = Font(name='Calibri')
zfont = Font(color = '006100')
zfill = PatternFill("solid", bgColor = "CCCCCC")
side = Side(border_style='thin', color="FF4682B4")
Border = Border(bottom=side, top=side)
titlefill = PatternFill('solid', fgColor = 'FF95B3D7')
titlefont = Font(name = 'Calibri', bold = True, italic = False, strike = False, color = "FF000000")
fillEukaryote = PatternFill('solid', fgColor = 'FF87CEFA')
fontEukaryote = Font(name = 'Calibri', bold = False, italic = False, strike = False, color = "FF9C0006")
fillVirus = PatternFill('solid', fgColor = 'FFEEEE00')
fontred = Font(name = 'Calibri', bold = False, italic = False, strike = False, color = "FF9C0006")
fontclinical = Font(name = 'Calibri', bold = False, italic = False, strike = False, color = '2F4F4F')
fillpathogen = PatternFill('solid', bgColor = "FFFFC7CE", fgColor = "FFFFC7CE")
fillclinical = PatternFill('solid', bgColor = "FFAFEEEE", fgColor = "FFAFEEEE")

wb = Workbook()
ws= wb.create_sheet("Pathoseq", 0)
rownum = 1
with open(sys.argv[1]+f'/mNGS.{date}.merge.KK.xls', 'r') as nameid:
    for line in nameid:
        linecontent = line.strip('\n').split("\t")
        species = linecontent[2]
        for i in range(0, len(linecontent)):
            if linecontent[i].isdigit():
                linecontent[i] = int(linecontent[i])
            cell = ws.cell(row = rownum, column = i+1, value = linecontent[i])
            cell.font = font
            cell.border = Border
            if rownum == 3:
                cell.font = titlefont
                cell.fill = titlefill
            if i == 0:
                if cell.value == "D_2759_Eukaryota":
                    cell.font = fontEukaryote
                    cell.fill = fillEukaryote
                elif cell.value == "D_10239_Viruses":
                    cell.font = fontEukaryote
                    cell.fill = fillVirus
            if i == len(linecontent)-1:
                if linecontent[i] == "clinical":
                    cell.fill = fillclinical
                    cell.font = fontclinical
                elif linecontent[i] == "pathogen":
                    cell.fill = fillpathogen
                    cell.font = fontred
        rownum += 1

data = f'A4:{get_column_letter(len(list(ws.columns)))}{len(list(ws.rows))}'
ws.conditional_formatting.add(data, CellIsRule(operator='==', formula=[0], font=zfont, fill=zfill))

wb.save(sys.argv[1]+f'/GZ.mNGS.{date}.mergeKK.xlsx')

#ddRobot.robot('广州 kraken2 运行完毕！', 'WHQ')

#url = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=90434457-2bc2-45f4-8c5d-4c781ec72314"
#data = {"msgtype": "text", "text": {"content": "广州 kraken2 运行完毕！",}}
#r = requests.post(url, json=data)
